import numpy as np
import random

class GridMap:
    def __init__(self, size=(100, 100), resolution=1.0):
        """
        Initialize a grid map for path planning simulation
        :param size: (width, height) of the grid in cells
        :param resolution: grid cell size in meters
        """
        self.width = size[0]
        self.height = size[1]
        self.resolution = resolution
        self.grid = np.zeros((self.height, self.width), dtype=np.int8)
        self.dynamic_obstacles = []

    def generate_random_obstacles(self, obstacle_density=0.15):
        """
        Generate random static obstacles in the grid
        :param obstacle_density: proportion of grid to be obstacles (0.0-1.0)
        """
        num_obstacles = int(self.width * self.height * obstacle_density)
        obstacle_coords = set()

        # Ensure start and goal positions are kept free
        start_area = [(0, 0), (0, 1), (1, 0), (1, 1)]
        goal_area = [(self.width-1, self.height-1), (self.width-2, self.height-1),
                    (self.width-1, self.height-2), (self.width-2, self.height-2)]
        reserved = set(start_area + goal_area)

        while len(obstacle_coords) < num_obstacles:
            x = random.randint(0, self.width-1)
            y = random.randint(0, self.height-1)
            if (x, y) not in reserved and (x, y) not in obstacle_coords:
                # Create obstacle clusters
                cluster_size = random.randint(1, 5)
                for i in range(cluster_size):
                    for j in range(cluster_size):
                        cx = x + i - cluster_size//2
                        cy = y + j - cluster_size//2
                        if (0 <= cx < self.width and 0 <= cy < self.height and
                            (cx, cy) not in reserved and (cx, cy) not in obstacle_coords):
                            obstacle_coords.add((cx, cy))
                            if len(obstacle_coords) >= num_obstacles:
                                break
                    if len(obstacle_coords) >= num_obstacles:
                        break

        # Update grid with obstacles
        for (x, y) in obstacle_coords:
            self.grid[y, x] = 1

        return self.grid

    def add_dynamic_obstacle(self, start_pos, end_pos, speed=0.5):
        """
        Add a dynamic obstacle with linear movement
        :param start_pos: (x, y) start position in world coordinates
        :param end_pos: (x, y) end position in world coordinates
        :param speed: m/s
        """
        start_grid = self.world2grid(start_pos)
        end_grid = self.world2grid(end_pos)

        obstacle = {
            'start': start_grid,
            'end': end_grid,
            'current': start_grid,
            'speed': speed,
            'direction': np.arctan2(end_grid[1]-start_grid[1], end_grid[0]-start_grid[0]),
            'path': self._generate_straight_path(start_grid, end_grid)
        }

        self.dynamic_obstacles.append(obstacle)
        return len(self.dynamic_obstacles) - 1

    def update_dynamic_obstacles(self, dt=0.1):
        """
        Update positions of all dynamic obstacles
        :param dt: time step in seconds
        """
        obstacle_positions = []

        for obs in self.dynamic_obstacles:
            # Calculate distance to move
            move_distance = obs['speed'] * dt / self.resolution
            obs['current'] = (obs['current'][0] + move_distance * np.cos(obs['direction']),
                             obs['current'][1] + move_distance * np.sin(obs['direction']))

            # Check if obstacle has reached end
            dist_to_end = np.hypot(obs['current'][0] - obs['end'][0],
                                  obs['current'][1] - obs['end'][1])

            if dist_to_end < move_distance:
                # Reverse direction
                obs['direction'] += np.pi
                obs['start'], obs['end'] = obs['end'], obs['start']

            # Add current position to obstacle positions
            obstacle_positions.append((int(round(obs['current'][0])), int(round(obs['current'][1]))))

        return obstacle_positions

    def is_occupied(self, x, y, check_dynamic=True):
        """
        Check if a grid cell is occupied by an obstacle
        :param x: grid x coordinate
        :param y: grid y coordinate
        :param check_dynamic: whether to check dynamic obstacles
        :return: True if occupied
        """
        if not (0 <= x < self.width and 0 <= y < self.height):
            return True  # Out of bounds is considered occupied

        if self.grid[y, x] == 1:
            return True  # Static obstacle

        if check_dynamic:
            # Check dynamic obstacles
            for obs in self.dynamic_obstacles:
                ox, oy = int(round(obs['current'][0])), int(round(obs['current'][1]))
                if (x, y) == (ox, oy):
                    return True

        return False

    def world2grid(self, world_coords):
        """Convert world coordinates to grid indices"""
        return [int(world_coords[0] / self.resolution), int(world_coords[1] / self.resolution)]

    def grid2world(self, grid_indices):
        """Convert grid indices to world coordinates"""
        return [grid_indices[0] * self.resolution, grid_indices[1] * self.resolution]

    def _generate_straight_path(self, start, end):
        """Generate straight path between two points"""
        path = []
        dx = end[0] - start[0]
        dy = end[1] - start[1]
        steps = max(abs(dx), abs(dy)) + 1

        for i in range(steps):
            x = start[0] + dx * i / steps
            y = start[1] + dy * i / steps
            path.append((x, y))

        return path

    def get_grid(self):
        """Return the current grid map"""
        return self.grid.copy()

    def reset(self):
        """Reset the grid map"""
        self.grid = np.zeros((self.height, self.width), dtype=np.int8)
        self.dynamic_obstacles = []